package org.iscas.tcse.preprocessing;

import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.ml.linalg.DenseVector;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.Transaction;

import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;

public class MinMaxScaler extends ProcessFunction<DenseVector, DenseVector>
        implements Preprocessing {
    private static final Logger LOG = LoggerFactory.getLogger(MinMaxScaler.class);

    public final String prefix;

    private final int windowSize;
    private int currSize = 0;
    private final int[] featureIndices;

    private volatile boolean bootWindow = true;

    private volatile double[] minValues;
    private volatile double[] maxValues;

    private volatile boolean isRunning = true;

    private final String[] minKeys;
    private final String[] maxKeys;

    private final BlockingQueue<Tuple2<double[], double[]>> channel;
    private final String host;
    private final int port;

    private final String setMinScript =
            "local c = tonumber(redis.call('get', KEYS[1])); if c then if tonumber(ARGV[1]) < c then redis.call('set', KEYS[1], ARGV[1]); return ARGV[1] else return tostring(c) end else redis.call('set', KEYS[1], ARGV[1]); return ARGV[1] end";
    private final String setMaxScript =
            "local c = tonumber(redis.call('get', KEYS[1])); if c then if tonumber(ARGV[1]) > c then redis.call('set', KEYS[1], ARGV[1]); return ARGV[1] else return tostring(c) end else redis.call('set', KEYS[1], ARGV[1]); return ARGV[1] end";

    private final String setMinScriptSha;
    private final String setMaxScriptSha;

    public MinMaxScaler(int windowSize, int[] featureIndices, String host, int port, String name) {
        this.windowSize = windowSize;
        this.featureIndices = featureIndices;
        this.host = host;
        this.port = port;
        this.prefix = name + "_min_max_scaler";

        minKeys = new String[featureIndices.length];
        maxKeys = new String[featureIndices.length];
        minValues = new double[featureIndices.length];
        maxValues = new double[featureIndices.length];
        channel = new ArrayBlockingQueue<>(1);

        Jedis jedis = new Jedis(host, port);

        Transaction tx = jedis.multi();

        for (int i = 0; i < featureIndices.length; i++) {
            minKeys[i] = getMinKey(featureIndices[i]);
            maxKeys[i] = getMaxKey(featureIndices[i]);
            minValues[i] = Double.MAX_VALUE;
            maxValues[i] = Double.MIN_VALUE;
            tx.del(minKeys[i]);
            tx.del(maxKeys[i]);
        }
        tx.exec();
        tx.close();

        this.setMinScriptSha = jedis.scriptLoad(setMinScript);
        this.setMaxScriptSha = jedis.scriptLoad(setMaxScript);

        jedis.close();
    }

    @Override
    public void close() throws Exception {
        super.close();
        channel.offer(new Tuple2<>(new double[0], new double[0]));
        isRunning = false;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);

        new Thread(
                        () -> {
                            try {
                                syncState();
                            } catch (Exception e) {
                                LOG.info("syncRoutine error: {}", e.getMessage());
                            }
                        })
                .start();
        //        LOG.info("opened");
    }

    @Override
    public void processElement(
            DenseVector value,
            ProcessFunction<DenseVector, DenseVector>.Context ctx,
            Collector<DenseVector> out)
            throws Exception {
        for (int i = 0; i < featureIndices.length; i++) {
            minValues[i] = Math.min(value.get(featureIndices[i]), minValues[i]);
            maxValues[i] = Math.max(value.get(featureIndices[i]), maxValues[i]);
        }
        currSize++;

        // process one record
        if (!bootWindow) {
            for (int i = 0; i < featureIndices.length; i++) {
                value.set(
                        featureIndices[i],
                        (value.get(featureIndices[i]) - minValues[i])
                                / (maxValues[i] - minValues[i]));
            }
        }
        out.collect(value);

        if (currSize >= windowSize) {
            double[] minValuesBuffer = this.minValues.clone();
            double[] maxValuesBuffer = this.maxValues.clone();
            currSize = 0;
            channel.put(new Tuple2<>(minValuesBuffer, maxValuesBuffer));
        }
    }

    private String getMinKey(int featureIdx) {
        return prefix + "_min_" + featureIdx;
    }

    private String getMaxKey(int featureIdx) {
        return prefix + "_max_" + featureIdx;
    }

    public void syncState() throws Exception {
        Jedis jedis = new Jedis(host, port);
        while (isRunning) {
            Tuple2<double[], double[]> data = channel.take();
            if (data.f0.length == 0) {
                //                LOG.info("syncRoutine exit");
                break;
            }
            Transaction tx = jedis.multi();
            double[] minValuesBuffer = data.f0;
            double[] maxValuesBuffer = data.f1;

            for (int i = 0; i < featureIndices.length; i++) {
                tx.evalsha(setMinScriptSha, 1, minKeys[i], String.valueOf(minValuesBuffer[i]));
                tx.evalsha(setMaxScriptSha, 1, maxKeys[i], String.valueOf(maxValuesBuffer[i]));
            }

            List<Object> result = tx.exec();

            boolean flag = true;
            for (int i = 0; i < featureIndices.length; i++) {
                Object minObj = result.get(i * 2);
                Object maxObj = result.get(i * 2 + 1);
                if (minObj == null || maxObj == null) {
                    flag = false;
                    break;
                }
                //                LOG.info("min: {}, max: {}", minObj, maxObj);
                minValuesBuffer[i] = Double.parseDouble((String) minObj);
                maxValuesBuffer[i] = Double.parseDouble((String) maxObj);
            }
            tx.close();
            if (!flag) {
                continue;
            }

            this.minValues = minValuesBuffer;
            this.maxValues = maxValuesBuffer;
            this.bootWindow = false;
        }
        jedis.close();
    }
}
